import numpy as np
import pandas as pd

class DataProcessor:
    def __init__(self, subchannel):
        self.subchannel = subchannel
        
    def calculate_average_void_fraction(self):
        """计算平均空泡率"""
        return np.mean(self.subchannel.void_fraction, axis=1)
        
    def calculate_channel_flow_rates(self):
        """计算各通道的流量"""
        flow_rates = np.zeros((self.subchannel.n_axial, self.subchannel.n_channels))
        
        for i in range(self.subchannel.n_axial):
            for j in range(self.subchannel.n_channels):
                rho_l = self.subchannel.fluid_prop.get_saturation_properties(
                    self.subchannel.pressure[i,j])[1]
                rho_v = self.subchannel.fluid_prop.get_saturation_properties(
                    self.subchannel.pressure[i,j])[2]
                alpha = self.subchannel.void_fraction[i,j]
                
                flow_rates[i,j] = (1-alpha) * rho_l * self.subchannel.vel_liquid[i,j] + \
                                 alpha * rho_v * self.subchannel.vel_vapor[i,j]
                                 
        return flow_rates
        
    def export_results_to_excel(self, filename):
        """导出结果到Excel文件"""
        # 创建一个ExcelWriter对象
        with pd.ExcelWriter(filename) as writer:
            # 导出空泡率数据
            void_fraction_df = pd.DataFrame(
                self.subchannel.void_fraction,
                columns=[f'Channel_{i+1}' for i in range(self.subchannel.n_channels)]
            )
            void_fraction_df.to_excel(writer, sheet_name='Void_Fraction')
            
            # 导出压力数据
            pressure_df = pd.DataFrame(
                self.subchannel.pressure,
                columns=[f'Channel_{i+1}' for i in range(self.subchannel.n_channels)]
            )
            pressure_df.to_excel(writer, sheet_name='Pressure')
            
            # 导出流量数据
            flow_rates = self.calculate_channel_flow_rates()
            flow_rates_df = pd.DataFrame(
                flow_rates,
                columns=[f'Channel_{i+1}' for i in range(self.subchannel.n_channels)]
            )
            flow_rates_df.to_excel(writer, sheet_name='Flow_Rates')
            
    def generate_summary_statistics(self):
        """生成统计摘要"""
        stats = {
            'void_fraction': {
                'mean': np.mean(self.subchannel.void_fraction),
                'std': np.std(self.subchannel.void_fraction),
                'max': np.max(self.subchannel.void_fraction),
                'min': np.min(self.subchannel.void_fraction)
            },
            'pressure': {
                'mean': np.mean(self.subchannel.pressure),
                'std': np.std(self.subchannel.pressure),
                'max': np.max(self.subchannel.pressure),
                'min': np.min(self.subchannel.pressure)
            },
            'flow_rates': {
                'mean': np.mean(self.calculate_channel_flow_rates()),
                'std': np.std(self.calculate_channel_flow_rates()),
                'max': np.max(self.calculate_channel_flow_rates()),
                'min': np.min(self.calculate_channel_flow_rates())
            }
        }
        return stats 